from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait  # 显示等待
from selenium.webdriver.support import expected_conditions as ec  # 等待条件
from selenium.webdriver.chrome.service import Service
import openpyxl

service = Service('E:/devtools/Python/Python312/chromedriver.exe')  # 替换为实际路径
driver = webdriver.Chrome(service=service)  # 创建浏览器对象


class TrainSpider(object):
    # 登录地址
    login_url = 'https://kyfw.12306.cn/otn/resources/login.html'
    # 个人中心网址
    profile_url = 'https://kyfw.12306.cn/otn/view/index.html'
    # 余票查询
    left_ticket = 'https://kyfw.12306.cn/otn/leftTicket/init?linktypeid=dc'

    def __init__(self, from_station, to_station, train_date):
        self.from_station = from_station
        self.to_station = to_station
        self.train_date = train_date
        self.station_code = self.init_station_code()

    # 登录
    def login(self):
        # 打开登录页面
        driver.get(self.login_url)
        WebDriverWait(driver=driver, timeout=1000).until(
            ec.url_to_be(self.profile_url)  # 等待 url 变为个人中心
        )
        print("登录成功")

    def search_ticket(self):
        # 打开查询余票网址
        driver.get(self.left_ticket)
        # 出发站
        from_staion_input = driver.find_element(By.ID, 'fromStation')
        # 到达站
        to_staion_input = driver.find_element(By.ID, 'toStation')
        # 出发日期
        train_date_input = driver.find_element(By.ID, 'train_date')
        from_staion_code = self.station_code[self.from_station]
        to_staion_code = self.station_code[self.to_station]

        # 执行js
        driver.execute_script('arguments[0].value="%s"' % from_staion_code, from_staion_input)
        driver.execute_script('arguments[0].value="%s"' % to_staion_code, to_staion_input)
        driver.execute_script('arguments[0].value="%s"' % self.train_date, train_date_input)

    # 初始化车站代码
    def init_station_code(self):
        wb = openpyxl.load_workbook('车站代码.xlsx')
        ws = wb.active
        lst = []
        for row in ws.rows:
            sub_list = []
            for cell in row:
                sub_list.append(cell.value)
            lst.append(sub_list)

        return dict(lst)

    # 负责调用其他方法
    def run(self):
        # 登录
        self.login()
        # 余票查询
        self.search_ticket()


def start():
    spider = TrainSpider('北京', '上海', '2025-05-21')
    spider.run()


if __name__ == '__main__':
    start()
